from pysr import PySRRegressor
import numpy as np
from sklearn.model_selection import train_test_split
from bishe_situations.utils import PYSR_PARAMS, BaselineTrainer, generate_data, UNARY_OPS_PYSR, BINARY_OPS_PYSR

class PysrTrainer(BaselineTrainer):
    def __init__(self, params: dict = PYSR_PARAMS, population_size: int = 50, generations: int = 20, **kwargs):
        super().__init__(params, population_size, generations, **kwargs)
        self.model = PySRRegressor(
            binary_operators=BINARY_OPS_PYSR,
            unary_operators=UNARY_OPS_PYSR,
            niterations=generations,
            population_size=population_size,
            maxsize=25,
            parsimony=0.1,
            progress=False,
            verbosity=1,
            constraints={'^': (-5, 5)},
            **params,
            **kwargs
        )

    def override_fit(self, X, y):
        X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)
        self.model.fit(X_train, y_train.ravel())
        self.mse = self.model.score(X_val, y_val.ravel())

    def best_function(self):
        return self.model.sympy()

def optimize_hyperparameters(X, y, n_trials=10):    
    # 初始化PySR模型
    model = PysrTrainer(generations=n_trials)
    
    # 训练模型
    model.fit(X, y)
    
    return model.best_function(), model.mse, model.train_time

def main():
    # 定义目标函数
    def target_function(X):
        return np.sin(X) * 2 + np.cos(X) - 80 + np.exp(X)
    
    # 生成数据
    X, y = generate_data(target_function, n_samples=512, noise=0.8)
    
    # 使用PySR进行拟合
    print("开始符号回归...")
    best_equation, mse, train_time = optimize_hyperparameters(X, y)
    
    print("\n最佳方程:")
    print(best_equation)
    print(f"验证集MSE: {mse:.4f}")
    print(f"训练时间: {train_time:.4f}秒")

if __name__ == "__main__":
    main() 